// Copyright 2022 jeff.li. and/or its affiliates.

#include <turbo/json/stringbuffer.h>
#include <turbo/json/writer.h>

#include <tbir/pipeline/pickle.h>
#include <tbir/runtime/container.h>
#include <tbir/runtime/container_private.h>
#include <tbir/runtime/ft_container.h>
#include <tbir/runtime/json_util.h>
#include <tbir/runtime/registry.h>

namespace tbir::runtime::pickle {

    static Unicode ToJsonStructStr(const Any &rtv) {
        rapidjson::Document doc;
        rapidjson::StringBuffer buffer;
        ToJsonStruct(rtv, doc, doc.GetAllocator());
        buffer.Clear();
        rapidjson::Writer<rapidjson::StringBuffer> writer(buffer);
        doc.Accept(writer);
        return String(buffer.GetString()).decode();
    }

    static RTValue FromJsonStructStr(const string_view &bytes) {
        rapidjson::Document doc;
        MXCHECK(JsonUtil::FromString(bytes, doc));
        return FromJsonStruct(doc);
    }

    TBIR_REGISTER_GLOBAL("runtime.pickle_ToJsonStruct").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 1) << "[runtime.pickle.ToJsonStruct] Expect 1 arguments but get "
                                  << args.size();
        return ToJsonStructStr(args[0]);
    });

    TBIR_REGISTER_GLOBAL("runtime.pickle_FromJsonStruct").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 1) << "[runtime.pickle.FromJsonStruct] Expect 1 arguments but get "
                                  << args.size();
        String v = UnicodeHelper::Encode(args[0].As<unicode_view>());
        return FromJsonStructStr(v);
    });

    TBIR_REGISTER_GLOBAL("runtime.Serialize").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 1) << "[runtime.pickle.Serialize] Expect 1 arguments but get "
                                  << args.size();
        return Serialize(args[0]).decode();
    });

    TBIR_REGISTER_GLOBAL("runtime.DeSerialize").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 1) << "[runtime.pickle.DeSerialize] Expect 1 arguments but get "
                                  << args.size();
        auto v = UnicodeHelper::Encode(args[0].As<unicode_view>());
        return DeSerialize(v);
    });

}  // namespace tbir::runtime::pickle
